package lib

import "algorithms/ds"

func init() {
	Probs = append(Probs, Problem{
		Num:         337,
		Discription: "二叉树打家劫舍",
		Level:       2,
		Labels: map[string]int{
			"二叉树":    1,
			"递归":     1,
			"回溯":     1,
			"有限状态机":     1,
			"树形动态规划": 1,
		},
	})
}

func Rob3(root *ds.TreeNode) int {
	if root == nil {
		return 0
	}

	//每个节点返回两种状态：偷该节点的最大收益，不偷该节点的最大收益
	var dfs func(*ds.TreeNode) (int, int)
	dfs = func(node *ds.TreeNode) (int, int) {
		if node == nil {
			return 0, 0
		}

		leftSteal, leftCooldown := dfs(node.Left)
		rightSteal, rightCooldown := dfs(node.Right)
		return leftCooldown + rightCooldown + node.Val, max(leftCooldown, leftSteal) + max(rightCooldown, rightSteal)
	}

	dfs(root)
	steal, cooldown := dfs(root)
	return max(steal, cooldown)
}
